package org.budo.dubbo.protocol.async.repository.kafka;

import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.utils.Bytes;
import org.budo.dubbo.protocol.async.repository.BudoAsyncInvocation;
import org.budo.dubbo.protocol.async.util.AsyncLogUtil;
import org.budo.dubbo.protocol.async.util.AsyncSerializeUtil;
import org.budo.dubbo.protocol.async.util.DestinationUtil;
import org.budo.graph.metric.MetricUtil;

import com.alibaba.dubbo.rpc.Invocation;
import com.alibaba.dubbo.rpc.Invoker;
import com.alibaba.dubbo.rpc.RpcInvocation;

import lombok.extern.slf4j.Slf4j;

/**
 * @author lmw
 */
@Slf4j
public class KafkaAsyncRepository extends AbstractKafkaAsyncRepository {
    @Override
    public void sendMessageWhenInvoke(Invocation invocation) {
        String destinationName = DestinationUtil.getDestinationName(invocation.getInvoker());

        BudoAsyncInvocation budoAsyncInvocation = BudoAsyncInvocation.fromDubboInvocation(invocation);
        byte[] requestBody = AsyncSerializeUtil.serializeAsyncInvocation(budoAsyncInvocation, invocation.getInvoker());
        this.sendRawMessage(destinationName, requestBody);
    }

    @Override
    public void sendMessage(String destinationName, Object message) {
        byte[] requestBody = AsyncSerializeUtil.serializeAsyncInvocation(message, null);

        this.sendRawMessage(destinationName, requestBody);
    }

    @Override
    public void sendRawMessage(String destinationName, Object rawMessage) {
        String key = "async.kafka." + destinationName + ".sendRawMessage";
        String uuid = MetricUtil.before(key, true); // 性能统计 1

        if (null == destinationName || destinationName.trim().isEmpty()) {
            MetricUtil.after(key, uuid); // 性能统计 2
            throw new IllegalArgumentException("#38 destinationName=" + destinationName + ", rawMessage=" + rawMessage);
        }

        KafkaProducer<String, Bytes> kafkaProducer = this.getKafkaProducer(destinationName);

        byte[] messageBytes;
        if (rawMessage instanceof byte[]) {
            messageBytes = (byte[]) rawMessage;
        } else if (rawMessage instanceof String) {
            messageBytes = ((String) rawMessage).getBytes();
        } else {
            MetricUtil.after(key, uuid); // 性能统计 3
            throw new IllegalArgumentException("#55 rawMessage.type=" + rawMessage.getClass());
        }

        // send
        ProducerRecord<String, Bytes> producerRecord = new ProducerRecord<String, Bytes>(destinationName, Bytes.wrap(messageBytes));
        kafkaProducer.send(producerRecord);

        log.debug("#50 sendMessageWhenInvoke, destinationName={}", destinationName);

        // 性能统计 4
        MetricUtil.after(key, uuid);
    }

    @Override
    public Integer popMessageAndInvoke(Invoker<?> invoker) {
        long start = System.currentTimeMillis();

        String destinationName = DestinationUtil.getDestinationName(invoker);

        // poll message
        KafkaConsumer<String, Bytes> kafkaConsumer = this.getKafkaConsumer(destinationName);
        ConsumerRecords<String, Bytes> consumerRecords = kafkaConsumer.poll(POLL_TIMEOUT);

        if (log.isDebugEnabled()) {
            log.debug("#82 poll cost " + (System.currentTimeMillis() - start) + ", consumerRecords.count=" + consumerRecords.count() //
                    + ", POLL_TIMEOUT=" + POLL_TIMEOUT + ", destinationName=" + destinationName);
        }

        // invoke
        int messageCount = 0;
        for (ConsumerRecord<String, Bytes> consumerRecord : consumerRecords) {
            String key = "async.kafka." + destinationName + ".popMessageAndInvoke";
            String uuid = MetricUtil.before(key, true); // 性能统计 开始

            byte[] messageBody = consumerRecord.value().get();
            BudoAsyncInvocation budoAsyncInvocation = AsyncSerializeUtil.deserializeAsyncInvocation(messageBody, invoker);

            if (log.isDebugEnabled()) {
                log.debug("#58 message received, " + consumerRecord + ", budoAsyncInvocation=" + budoAsyncInvocation);
            }

            // 如果消息延迟就打日志
            AsyncLogUtil.logMessageDelay(budoAsyncInvocation, destinationName, invoker, this);

            RpcInvocation invocation = BudoAsyncInvocation.toDubboInvocation(budoAsyncInvocation);
            invoker.invoke(invocation); // 执行方法
            messageCount++;

            MetricUtil.after(key, uuid); // 性能统计 结束
        }

        // 方法调用成功后，提交确认这一批消息
        this.commitOffsetSync(kafkaConsumer, messageCount);

        return messageCount;
    }

    @Override
    public String toString() {
        return super.toString() + ", properties=" + this.getProperties();
    }
}